import requests
import numpy as np
from typing import List, Dict, Any
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class HyDERetriever:
    """
    HyDE (Hypothetical Document Embeddings) 检索器

    该类实现了完整的HyDE流程：
    1. 使用LLM生成假设文档
    2. 对假设文档进行向量嵌入
    3. 在知识库中检索相似文档
    4. 结合检索结果生成最终答案
    """

    def __init__(self,
                 model_url: str = "http://localhost:11434/api/generate",
                 model_name: str = "deepseek-r1:1.5b",
                 embedding_model: str = "all-MiniLM-L6-v2"):
        """
        初始化HyDE检索器

        Args:
            model_url: DeepSeek模型API地址
            model_name: 模型名称
            embedding_model: 向量嵌入模型名称
        """
        self.model_url = model_url
        self.model_name = model_name

        # 初始化向量嵌入模型
        logger.info(f"加载嵌入模型: {embedding_model}")
        self.embedding_model = SentenceTransformer(embedding_model)

        # 知识库存储
        self.knowledge_base = []
        self.knowledge_embeddings = None

    def add_documents(self, documents: List[str]):
        """
        添加文档到知识库

        Args:
            documents: 文档列表
        """
        logger.info(f"添加 {len(documents)} 个文档到知识库")
        self.knowledge_base.extend(documents)

        # 计算所有文档的向量嵌入
        all_docs = self.knowledge_base
        self.knowledge_embeddings = self.embedding_model.encode(all_docs)
        logger.info("文档向量嵌入计算完成")

    def generate_hypothetical_document(self, query: str) -> str:
        """
        使用DeepSeek模型生成假设文档

        Args:
            query: 用户查询

        Returns:
            生成的假设文档
        """
        logger.info(f"为查询生成假设文档: {query}")

        prompt = f"""基于以下问题，请生成一个详细的假设性回答文档。这个文档应该包含可能的答案内容、相关技术细节和背景信息。

问题: {query}

请生成一个结构清晰、内容详细的假设文档："""

        payload = {
            "model": self.model_name,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": 0.7,
                "max_tokens": 512
            }
        }

        try:
            response = requests.post(self.model_url, json=payload, timeout=30)
            response.raise_for_status()

            result = response.json()
            hypothetical_doc = result.get('response', '').strip()

            logger.info("假设文档生成成功")
            return hypothetical_doc

        except requests.exceptions.RequestException as e:
            logger.error(f"请求DeepSeek模型时出错: {e}")
            # 备用方案：返回简单的假设文档
            return f"关于{query}的相关信息和解答内容。"

    def retrieve_similar_documents(self, hypothetical_doc: str, top_k: int = 3) -> List[Dict[str, Any]]:
        """
        基于假设文档检索相似的真实文档

        Args:
            hypothetical_doc: 假设文档
            top_k: 返回最相似的文档数量

        Returns:
            相似文档列表，包含文档内容和相似度分数
        """
        if self.knowledge_embeddings is None:
            logger.warning("知识库为空，无法进行检索")
            return []

        logger.info("开始检索相似文档")

        # 计算假设文档的向量嵌入
        hyp_embedding = self.embedding_model.encode([hypothetical_doc])

        # 计算与知识库中所有文档的相似度
        similarities = cosine_similarity(hyp_embedding, self.knowledge_embeddings)[0]

        # 获取最相似的文档索引
        top_indices = np.argsort(similarities)[::-1][:top_k]

        similar_docs = []
        for idx in top_indices:
            similar_docs.append({
                'content': self.knowledge_base[idx],
                'similarity': float(similarities[idx]),
                'index': int(idx)
            })

        logger.info(f"检索到 {len(similar_docs)} 个相似文档")
        return similar_docs

    def generate_final_answer(self, query: str, hypothetical_doc: str,
                              retrieved_docs: List[Dict[str, Any]]) -> str:
        """
        基于假设文档和检索文档生成最终答案

        Args:
            query: 原始查询
            hypothetical_doc: 假设文档
            retrieved_docs: 检索到的相似文档

        Returns:
            最终生成的答案
        """
        logger.info("生成最终答案")

        # 构建上下文
        context = "参考文档:\n"
        for i, doc in enumerate(retrieved_docs, 1):
            context += f"{i}. {doc['content']}\n"

        prompt = f"""基于以下信息回答用户问题：

用户问题: {query}

假设文档: {hypothetical_doc}

{context}

请结合假设文档和参考文档，给出准确、详细的回答："""

        payload = {
            "model": self.model_name,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": 0.3,
                "max_tokens": 1024
            }
        }

        try:
            response = requests.post(self.model_url, json=payload, timeout=30)
            response.raise_for_status()

            result = response.json()
            final_answer = result.get('response', '').strip()

            logger.info("最终答案生成成功")
            return final_answer

        except requests.exceptions.RequestException as e:
            logger.error(f"生成最终答案时出错: {e}")
            return "抱歉，生成答案时遇到问题。"

    def hyde_retrieve(self, query: str, top_k: int = 3) -> Dict[str, Any]:
        """
        执行完整的HyDE检索流程

        Args:
            query: 用户查询
            top_k: 检索的文档数量

        Returns:
            包含假设文档、检索文档和最终答案的字典
        """
        logger.info(f"开始HyDE检索流程: {query}")

        # 步骤1: 生成假设文档
        hypothetical_doc = self.generate_hypothetical_document(query)

        # 步骤2: 检索相似文档
        retrieved_docs = self.retrieve_similar_documents(hypothetical_doc, top_k)

        # 步骤3: 生成最终答案
        final_answer = self.generate_final_answer(query, hypothetical_doc, retrieved_docs)

        return {
            'query': query,
            'hypothetical_document': hypothetical_doc,
            'retrieved_documents': retrieved_docs,
            'final_answer': final_answer
        }